#!/bin/bash
# =================================================================
# 使用wandb和预训练模型训练MOTIP (8-GPU)
# 建议在 tmux 或类似工具中运行此脚本
# =================================================================

# --- 配置 ---

# 设置项目根目录（MA-MOTIP）
PROJECT_DIR="/home/liuyonghui/code/MA-MOTIP"
cd "$PROJECT_DIR" || exit

# 输出目录
# 训练产生的模型和日志将保存在此
OUTPUTS_DIR="${OUTPUTS_DIR:-${PROJECT_DIR}/outputs/motip_dancetrack_wandb}"

# 实验名称
# 这将用作 wandb run 的名称
EXP_NAME="${EXP_NAME:-motip_dancetrack_lwg}"

# Wandb 配置
# 请替换为您的 wandb entity 和 project 名称
export WANDB_ENTITY="yves-det"
export WANDB_PROJECT="MOTIP"
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export WANDB_BASE_URL=https://api.bandw.top
# 数据集根目录
# 请确保路径指向包含 DanceTrack 数据集的文件夹
DATA_ROOT="${DATA_ROOT:-/home/liuyonghui/datasets}"

# 预训练模型路径
# 根据 GET_STARTED.md，预训练模型应放置在 ./pretrains/ 目录下
# DanceTrack 预训练模型: r50_deformable_detr_coco_dancetrack.pth
PRETRAIN_PATH="${PRETRAIN_PATH:-${PROJECT_DIR}/pretrains/r50_deformable_detr_coco_dancetrack.pth}"

# 配置文件
# 使用带 LWG + Buffer Gate 的配置
CONFIG_PATH="${CONFIG_PATH:-${PROJECT_DIR}/configs/r50_deformable_detr_motip_dancetrack_lwg.yaml}"

# LWG / Buffer Gate 相关参数（可通过环境变量覆盖）
LWG_MODEL_PATH="${LWG_MODEL_PATH:-${PROJECT_DIR}/pretrains/lwg_dancetrack.pt}"
BUFFER_GATE_MODEL_PATH="${BUFFER_GATE_MODEL_PATH:-${PROJECT_DIR}/pretrains/buffer_gate_dancetrack.pt}"
LWG_MAIN_TH="${LWG_MAIN_TH:-0.60}"
LWG_BUF_TH="${LWG_BUF_TH:-0.40}"
BUFFER_CAPACITY="${BUFFER_CAPACITY:-10}"
BUFFER_PROMOTE_STEPS="${BUFFER_PROMOTE_STEPS:-3}"
BUFFER_GATE_THRESH="${BUFFER_GATE_THRESH:-0.50}"

# GPU 数量
NUM_GPUS="${NUM_GPUS:-8}"

# --- 环境设置 ---

# 避免多进程刷屏
export NVI_NOTIFY_IGNORE_TASK=1
# 解决AMD平台MKL库性能问题
export LD_PRELOAD=/opt/lib/fakeintel/libfakeintel.so

# 如果您使用特定的conda环境，请取消下面的注释并修改环境名称
echo "Activating conda environment..."
source /opt/miniconda3/etc/profile.d/conda.sh
conda activate MOTIP

# 准备日志文件
LOG_DIR="$PROJECT_DIR/logs"
mkdir -p "$LOG_DIR"
LOG_FILE="$LOG_DIR/${EXP_NAME}_$(date +%Y%m%d_%H%M%S).log"

echo "项目目录: $PROJECT_DIR"
echo "数据集目录: $DATA_ROOT"
echo "输出目录: $OUTPUTS_DIR"
echo "日志文件: $LOG_FILE"
echo "GPU数量: $NUM_GPUS"
echo "Wandb Entity: $WANDB_ENTITY"
echo "Wandb Project: $WANDB_PROJECT"

# --- 检查文件是否存在 ---

if [ ! -f "$PRETRAIN_PATH" ]; then
    echo "错误: 预训练模型未找到,路径: $PRETRAIN_PATH"
    echo "请从 Model Zoo 下载并放置在正确的位置。"
    exit 1
fi

if [ ! -d "$DATA_ROOT" ]; then
    echo "错误: 数据集目录未找到,路径: $DATA_ROOT"
    exit 1
fi

# ==============================
# 启动分布式训练
# ==============================

echo "开始训练..."
accelerate launch --num_processes="$NUM_GPUS" train.py \
    --config-path "$CONFIG_PATH" \
    --data-root "$DATA_ROOT" \
    --detr-pretrain "$PRETRAIN_PATH" \
    --outputs-dir "$OUTPUTS_DIR" \
    --exp-name "$EXP_NAME" \
    --use-wandb True \
    --use-lwg True \
    --lwg-model-path "$LWG_MODEL_PATH" \
    --lwg-main-th "$LWG_MAIN_TH" \
    --lwg-buf-th "$LWG_BUF_TH" \
    --buffer-capacity "$BUFFER_CAPACITY" \
    --buffer-promote-steps "$BUFFER_PROMOTE_STEPS" \
    2>&1 | tee "$LOG_FILE"

echo "训练完成。日志已保存到 $LOG_FILE"
